﻿#include <iostream>
#include <torch/script.h>
#include <opencv2/opencv.hpp>


using namespace std;

cv::Mat load_img(const string& img_path){
  auto image = cv::imread(img_path,cv::ImreadModes::IMREAD_COLOR);
  cv::Mat image_transfomed;
  cv::resize(image, image_transfomed, cv::Size(640,640));
  cv::cvtColor(image_transfomed, image_transfomed, cv::COLOR_BGR2RGB);
  return image_transfomed;
}

torch::Tensor process_img(cv::Mat img){
  // 图像转换为Tensor
  torch::Tensor tensor_image = torch::from_blob(img.data, {img.rows, img.cols,3},torch::kByte);
  tensor_image = tensor_image.permute({2,0,1});
  tensor_image = tensor_image.toType(torch::kFloat32);
  tensor_image = tensor_image.div(255);
  tensor_image = tensor_image.unsqueeze(0);
  return tensor_image;
}


int main() {
  torch::jit::script::Module module;
  try {
    module = torch::jit::load("yolov8n.torchscript");
    module.to(torch::kCUDA, torch::kFloat32);
  } catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }
  auto image_transfomed = load_img("bus.jpg");
  auto tensor_image = process_img(image_transfomed);

  at::Tensor output = module.forward({tensor_image}).toTensor();
  std::cout << torch::argmax(output)<< '\n';
  return 0;
}
